import cv2
import numpy as np
from skimage.metrics import structural_similarity as ssim


class StateAnalyzer:
    def __init__(self, template_dir='templates/'):
        self.templates = self._load_templates(template_dir)

    def _load_templates(self, path):
        """加载预存模板特征"""
        # 示例模板结构：{'sun_number': {'image': ndarray, 'threshold': 0.9}}
        return {...}

    def analyze_plant_selection(self, screen_img):
        """解析选植物界面状态"""
        state = {'scene_type': 'plant_selection'}

        # 阳光识别
        sun_roi = self._get_dynamic_roi(screen_img, 'sun_area')
        state['initial_sun'] = self._recognize_sun(sun_roi)

        # 僵尸类型识别
        zombie_roi = self._get_dynamic_roi(screen_img, 'zombie_list')
        state['zombie_types'] = self._detect_zombies(zombie_roi)

        # 植物槽位分析
        state.update(self._analyze_plant_slots(screen_img))
        return state

    def _recognize_sun(self, roi_image):
        """混合识别阳光数值"""
        # 使用Tesseract OCR与模板匹配双校验
        ocr_result = pytesseract.image_to_string(roi_image)
        template_match = self._template_match(roi_image, 'sun_digits')
        return self._vote_result(ocr_result, template_match)

    def _detect_zombies(self, roi_image):
        """基于YOLO的僵尸类型检测"""
        # 加载预训练模型
        model = cv2.dnn.readNet('zombie_yolo.weights', 'zombie_yolo.cfg')
        blob = cv2.dnn.blobFromImage(roi_image, 1 / 255, (416, 416))
        model.setInput(blob)
        outputs = model.forward(model.getUnconnectedOutLayersNames())
        # 解析检测结果...
        return ['普通僵尸', '路障僵尸']

    def _analyze_plant_slots(self, screen_img):
        """植物槽位状态分析"""
        # 使用SIFT特征匹配植物图标
        sift = cv2.SIFT_create()
        kp1, des1 = sift.detectAndCompute(screen_img, None)

        slot_states = {'available': [], 'selected': []}
        for plant_type, template in self.templates.items():
            kp2, des2 = template['keypoints'], template['descriptors']
            # 特征匹配
            matches = cv2.BFMatcher().knnMatch(des1, des2, k=2)
            good = [m for m, n in matches if m.distance < 0.75 * n.distance]
            if len(good) > template['min_matches']:
                # 判断是否在可选区域
                is_available = self._check_availability(screen_img, plant_type)
                slot_states['available'].append({
                    'plant_type': plant_type,
                    'is_available': is_available
                })
        return slot_states